/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */

#include "BTreeIndex.h"
#include "BTreeNode.h"
#include <iostream> // ADDED

using namespace std;

#define CS143_TRY(COMMAND)              \
do {                                    \
    RC retVal = COMMAND;                \
	if (retVal != 0) { return retVal; } \
} while (0)

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid    = -1;
    treeHeight = 0;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
	CS143_TRY(pf.open(indexname, mode));

	// create brand new pagefile
	if (pf.endPid() == 0 || pf.endPid() == 1) {
		rootPid = 1;
		treeHeight = 0;
		writeMetadata();
	}

	// pagefile has been created, and nodes exist
	else if (pf.endPid() > 1) {
		readMetadata();
	}

	return 0;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
	writeMetadata();
	return pf.close();
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
	PageId pid;

	// if nodes exist
	if (treeHeight > 0) {
		CS143_TRY(visitLeaf(key)); // traverse down tree and populate visitedPids vector
	}

	// if no nodes exist, initialize metadata page
	// => first time inserting
	else {
		treeHeight = 1;
		pid = rootPid;

		// need to actually write to page so that .endPid() can be used
		writeMetadata();

		BTLeafNode leaf;

		CS143_TRY(leaf.insert(key, rid));
		CS143_TRY(leaf.write(pid, pf));

		return 0;
	}

	// [visitedPids vector should be updated by now]

	pid = visitedPids.back();

	// cursor should be at bottom (leaf-level) now
	BTLeafNode leaf;
	CS143_TRY(leaf.read(pid, pf));

	// successfully inserted into leaf-node
	if (leaf.insert(key, rid) == 0) {
		CS143_TRY(leaf.write(pid, pf));
	}

	// leaf-node is full
	// => need to split
	else {
		BTLeafNode sibleaf;
		int sibkey;

		CS143_TRY(leaf.insertAndSplit(key, rid, sibleaf, sibkey));
		CS143_TRY(sibleaf.setNextNodePtr(leaf.getNextNodePtr()));
		CS143_TRY(sibleaf.write(pf.endPid(), pf));
		CS143_TRY(leaf.setNextNodePtr(pf.endPid()-1));
		CS143_TRY(leaf.write(pid, pf));

		if (treeHeight > 1) {
			CS143_TRY(insertUp(sibkey, leaf.getNextNodePtr())); // bubble upwards
		}
		else {
			BTNonLeafNode newRoot;
			newRoot.initializeRoot(pid, sibkey, leaf.getNextNodePtr());
			CS143_TRY(newRoot.write(pf.endPid(), pf));

			treeHeight++;
			rootPid = pf.endPid() - 1;
			writeMetadata();
		}
	}

	writeMetadata();
    return 0;
}

// after child is insertAndSplit'ed, call insertUp()
// the node to insert into (parentPid) is the parent of the caller (childPid)
RC BTreeIndex::insertUp(int key, PageId childPid) {

	// get the parent's pid to insert into
	visitedPids.pop_back();
	PageId parentPid = visitedPids.back();

	BTNonLeafNode node;
	CS143_TRY(node.read(parentPid, pf));

	// inserted successfully into parent
	if (node.insert(key, childPid) == 0) {
		CS143_TRY(node.write(parentPid, pf));

		return 0;
	}

	// parent is full, so insertAndSplit and insertUp again
	else {
		BTNonLeafNode sibnode;
		int midkey;

		CS143_TRY(node.insertAndSplit(key, childPid, sibnode, midkey));
		CS143_TRY(node.write(parentPid, pf));

		PageId unclePid = pf.endPid();
		CS143_TRY(sibnode.write(unclePid, pf));

		// make new root
		if (visitedPids.size() == 1) {
			BTNonLeafNode newRoot;

			CS143_TRY(newRoot.initializeRoot(parentPid, midkey, unclePid));
			CS143_TRY(newRoot.write(pf.endPid(), pf));

			treeHeight++;
			rootPid = pf.endPid() - 1;

			writeMetadata();

			return 0;
		}

		// splitting non-root
		else {
			CS143_TRY(insertUp(midkey, pf.endPid()-1));
		}
	}

	return 0;
}

/*
 * Find the leaf-node index entry whose key value is larger than or
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
	if (treeHeight < 1)
		return RC_NO_SUCH_RECORD;

	PageId pid = rootPid; // start at root

	// only a root node exists (of type BTLeafNode)
	if (treeHeight == 1) {
		BTLeafNode leaf;
		int eid;

		// read data from root node
		CS143_TRY(leaf.read(pid, pf));
		CS143_TRY(leaf.locate(searchKey, eid));

		cursor.pid = pid;
		cursor.eid = eid;
	}

	// multiple levels exist
	else {
		int currentLevel = 1;
		while (currentLevel < treeHeight) {
			BTNonLeafNode node;

			CS143_TRY(node.read(pid, pf));
			CS143_TRY(node.locateChildPtr(searchKey, pid));

			currentLevel++;
		}

		// if reached here, then we should be at the leaf-node level
		BTLeafNode leaf;
		int eid;

		CS143_TRY(leaf.read(pid, pf));
		CS143_TRY(leaf.locate(searchKey, eid));

		cursor.pid = pid;
		cursor.eid = eid;
	}

    return 0;
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move forward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
	if (cursor.pid < -1 || cursor.eid < 0)
		return RC_INVALID_CURSOR;

	if (cursor.pid >= pf.endPid() || cursor.pid == -1)
		return RC_END_OF_TREE;

    BTLeafNode leaf;
    CS143_TRY(leaf.read(cursor.pid, pf));
    CS143_TRY(leaf.readEntry(cursor.eid, key, rid));

    // if the next eid resides in the current page, then simply increment eid
    if (cursor.eid + 1 < leaf.getKeyCount()) {
		cursor.eid++;
    }

    // otherwise, next eid is in next page, if it exists
    else {
		cursor.pid = leaf.getNextNodePtr();
		cursor.eid = 0;
    }

    return 0;
}

void BTreeIndex::readMetadata() {
	char buffer[PageFile::PAGE_SIZE];
	pf.read(0, buffer);
	memcpy(&rootPid, buffer, sizeof(int));
	memcpy(&treeHeight, buffer + sizeof(int), sizeof(int));
}

void BTreeIndex::writeMetadata() {
	char buffer[PageFile::PAGE_SIZE];
	memset(buffer, 0, PageFile::PAGE_SIZE);
	memcpy(buffer, &rootPid, sizeof(int));
	memcpy(buffer + sizeof(int), &treeHeight, sizeof(int));
	pf.write(0, buffer);
}

RC BTreeIndex::visitLeaf(int searchKey) {

	visitedPids.clear();

	if (treeHeight == 1) {
		visitedPids.push_back(rootPid);
		return 0;
	}

	else {
		int currentLevel = 1;
		PageId pid = rootPid;

		// traversing down, beginning at root
		while (currentLevel < treeHeight) {
			BTNonLeafNode node;

			visitedPids.push_back(pid);

			CS143_TRY(node.read(pid, pf));
			CS143_TRY(node.locateChildPtr(searchKey, pid));

			currentLevel++;
		}

		// reached leaf node
		visitedPids.push_back(pid);
	}

	return 0;
}
